import base64
import math

from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5


def generate_key():
    rsa = RSA.generate(1024)
    private_key = rsa.exportKey()
    publick_key = rsa.publickey().exportKey()
    return private_key.decode(), publick_key.decode()


def rsa_encrypt(public_key: str, plain_data: str) -> list[str]:
    """
    rsa加密函数

    :params publick_key: 公钥
    :params others: 需要加密的信息（支持长密文，每96分段）
    :return: 加密后的密文（切分成列表）
    """
    default_length = 96
    content = []

    public_key = RSA.import_key(public_key)
    cipher = PKCS1_v1_5.new(public_key)
    for i in range(0, math.ceil(len(plain_data) / default_length)):
        temp = plain_data[i * default_length: (i + 1) * default_length]
        content.append(base64.b64encode(cipher.encrypt(temp.encode())).decode('utf-8'))
    return content


def rsa_decrypt(private_key: str, encrypted_data: list[str]) -> str:
    """
    rsa解密函数

    :prams private_key: 私钥
    :params others: 加密后的密文切分列表
    :return: 解密后原始信息
    """
    private_key = RSA.import_key(private_key)
    cipher = PKCS1_v1_5.new(private_key)
    result = ''
    for item in encrypted_data:
        #print(base64.b64decode(item))
        result += cipher.decrypt(base64.b64decode(item), 0).decode('utf-8')
    return result


if __name__ == "__main__":
    # pri, pub = generate_key() 可以这样用
    pri = """-----BEGIN RSA PRIVATE KEY-----
MIICXQIBAAKBgQChMkifsa7UPkkFHFdNvfGYV49bjdIcJfWIHXqpO5+xPD5GC8/+
AzryyoogDj5PFO4AFGO1m3H/ZhF8eYF5hpOes7K5qqaShQGqQk1ZeJnK3hrrtMAP
yH7PJOOeyq7jxL68dVxXISMN5AxG1MNBJ5UalwrZABfoqbINYKtr/5/+WQIDAQAB
AoGBAJjxIbDLIGrmUdfUhQ0RDDoT0cP3MoRci7LC7Y/GL4U0+IzzRjvnCrYz74vs
l7fmzs4qilRUEXxWdn7Ytjz8r+/gFwHs/CBrOvOWfiA5FKT61jy7OveFS/GXP6Cp
5XgrCrau5J8lrp0tysHaAppO8x477VmHlZXDxnHQdZFT5NSFAkEA3LHHk7p09Gqi
uNkA9SckvVgzUqSV8WY8i2Bv/YDqotGLI5i1XivATxj5lZu41xzP+CijaRTr6iQH
k+5nhG3E9wJBALr72Ua9wIwzH2SoZOrMOxRb+lPfN1uL75opekKTXu6K5myfUMdX
6MqkAFVnrrYhKZGunPCF9urGOccGx17Zky8CQQCFNs6YTnlgO6njWCGq2SjlHgNE
xFI595NeMiUHaTS4eTG6l+gKYvBfDl1BQFEPj7W9PpJ80IMivwLP+Z0Eo9m7AkBJ
8xbASoVfLUzKI2seDT2EVMgQ6aPcGeDd6G+I6v3Ne7Ob2/o3jIXttqTHIiVPtJfb
HOLGoUTBHzp0gzLQSV5TAkAo5UUdZvsQj547rJPruWUKq96N5/rp3wGsnqLund7C
FhDisKhY2DxYsx9veRIpxTED+TVJaR+MdwOqxzXoTWei
-----END RSA PRIVATE KEY-----"""
    pub = """-----BEGIN PUBLIC KEY-----
MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQChMkifsa7UPkkFHFdNvfGYV49b
jdIcJfWIHXqpO5+xPD5GC8/+AzryyoogDj5PFO4AFGO1m3H/ZhF8eYF5hpOes7K5
qqaShQGqQk1ZeJnK3hrrtMAPyH7PJOOeyq7jxL68dVxXISMN5AxG1MNBJ5UalwrZ
ABfoqbINYKtr/5/+WQIDAQAB
-----END PUBLIC KEY-----"""

    data = '{"username":"hello","password":"hello","password_confirm":"","captcha":"1234","email":"","phone":"","setting":false,"AES_key":"hello_next_hello"}'
    #print(f'[main]明文长度：{len(data)}')  # 最大长度117，大于这个值会报错，需要切分
    E = rsa_encrypt(pub, data)
    D = rsa_decrypt(pri, E)
    #print(f'[main]加密长度：[{len(E)}]{E}')
    #print(f'[main]解密文本：[{len(D)}]{D}')


